# -*- coding: utf-8 -*-
import re

import time
from sqlalchemy import create_engine, Column,String,Integer
from sqlalchemy.dialects.mysql import MEDIUMTEXT
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
import pymysql
from bs4 import BeautifulSoup as bs
from dingdian.models import VulNvd
from dingdian.init_mysql import session
Base = declarative_base()

import sys
reload(sys)
sys.setdefaultencoding('utf-8')

class VulNvd2(Base):
    # 表的名字:
    __tablename__ = 'vul_nvd'

    # 表的结构:
    id = Column(Integer, primary_key=True)
    text = Column(MEDIUMTEXT)
    url = Column(String(80),unique=True)

def saveHtml(file_name, file_content):
    #    注意windows文件命名的禁用符，比如 /
    with open(file_name.replace('/', '_') + ".html", "wb") as f:
        #   写文件用bytes而不是str，所以要转码
        f.write(file_content)

engine1 = create_engine('mysql+pymysql://root:950314@localhost:3306/vul_tmp?charset=utf8')
Base.metadata.create_all(engine1)
DBSession = sessionmaker(bind=engine1)
session1 = DBSession()




pattern = re.compile(r'CVE-\d{4}-\d+')


rows = session1.query(VulNvd2).count()
for row in range(1,rows+1):
    try:
        #print row
        htmls = session1.query(VulNvd2).filter(VulNvd2.id == row).all()
        url = htmls[0].url
        html = htmls[0].text.decode('string_escape')

        cveId = re.search(pattern, html).group()

        title = cveId
        publishedData = bs(html, 'lxml').find(attrs={'data-testid': "vuln-published-on"}).get_text()
        modifiedData = bs(html, 'lxml').find(attrs={'data-testid': "vuln-last-modified-on"}).get_text()

        publishedData=time.strftime("%Y-%m-%d",time.strptime(publishedData, "%m/%d/%Y"))
        modifiedData = time.strftime("%Y-%m-%d", time.strptime(modifiedData, "%m/%d/%Y"))

        description = bs(html, 'lxml').find(attrs={'data-testid': "vuln-description"}).get_text()
        vendorPatchs = bs(html, 'lxml').find("table", attrs={'data-testid': "vuln-hyperlinks-table"}).find_all("a")

        vendorPatch = ""
        for patch in vendorPatchs:
            vendorPatch = vendorPatch + patch.get_text() + "\r\n"
        #print vendorPatch
        type = bs(html, 'lxml').find(attrs={'data-testid': "vuln-technical-details-0-link"}).get_text() if bs(html, 'lxml').find(attrs={'data-testid': "vuln-technical-details-0-link"})!=None else None

        new_vul = VulNvd(cveID=cveId,publishTime=publishedData,updateTime=modifiedData,
                         url=url,title=title,type=type,
                         vendorPatch=vendorPatch,message=description)
        session.add(new_vul)
        session.commit()
    except Exception,e :
        print cveId+"\n",e

#htmls =session1.query(VulNvd2).filter(VulNvd2.id==3).all()
#html = htmls[0].text.decode('string_escape')

#new_vul = VulNvd(url=item['url'])
#            session.add(new_vul)
#            session.commit()
